#include <stdio.h>
#include <math.h>

#include "matMul_kernel.cu"

float* fillMatrix(const int, const int, const float);
float* padMatrix(float*, const int, const int, const int);
void displayMatrix(const float*, const int, const int);

int tileWidth;

//Globals
float *A_GPU, *B_GPU, *C_GPU; //GPU pointers

dim3 dimBlock, dimGrid; //Block and Grid Dimensions

int main()
{

	cudaError_t err = cudaSuccess;		
	err = cudaSetDevice(1);	
	if (err == cudaSuccess) { printf("\nNo problems setting to device 1!\n");}

	tileWidth = 2;	
	dimBlock.x = tileWidth;
	dimBlock.y = tileWidth;
	dimBlock.z = 1;

	int ax = 3;	
	int ay = 2;

	int bx = 2;
	int by = 2;
	
	int cx = ax;
	int cy = by;



	float* A = (float *)malloc(ax*ay*sizeof(float));
	float* B = (float *)malloc(bx*by*sizeof(float));
	float* C = (float *)malloc(cx*cy*sizeof(float));

	float* _A;
	float* _B;
	float* _C;	
	
	// Setup padded matrices
	int axDim = ax;
	int ayDim = ay;
	int bxDim = bx;
	int byDim = by;
	int cxDim = ax;
	int cyDim = by;
	
	if (ax%2 != 0) { cxDim = axDim = ax - (ax%2) + 2; }
	if (ay%2 != 0) { ayDim = ay - (ay%2) + 2; }
	if (bx%2 != 0) { bxDim = bx - (bx%2) + 2; }
	if (by%2 != 0) { cyDim = byDim = by - (by%2) + 2; }
	
	dimGrid.x = cxDim;//ceil((float)(ayDim)/tileWidth);
	dimGrid.y = cyDim;//ceil((float)(ayDim)/tileWidth);
	dimGrid.z = 1;



	err = cudaMalloc(&A_GPU, axDim*ayDim*sizeof(float));
	err = cudaMalloc(&B_GPU, bxDim*byDim*sizeof(float));
	err = cudaMalloc(&C_GPU, cxDim*cyDim*sizeof(float));

	if (err == cudaSuccess) { printf("\nNo problems allocating device memory!\n");}

	A = fillMatrix(ax,ay,2.0f);
	B = fillMatrix(bx,by,2.0f);
	C = fillMatrix(cx,cy,0.0f);

	_A = padMatrix(A, ax, ay, 2);
	_B = padMatrix(B, bx, by, 2);
	_C = padMatrix(C, cx, cy, 2);

	err = cudaMemcpy(A_GPU, _A, axDim*ayDim*sizeof(float), cudaMemcpyHostToDevice);
	err = cudaMemcpy(B_GPU, _B, bxDim*byDim*sizeof(float), cudaMemcpyHostToDevice);
	err = cudaMemcpy(C_GPU, _C, cxDim*cyDim*sizeof(float), cudaMemcpyHostToDevice);

	if (err == cudaSuccess) { printf("\nNo problems copying arrays to device!\n");}	

	MatrixMultiplySharedMemory<<<dimGrid, dimBlock, dimBlock.x*dimBlock.x*sizeof(float)>>>(A_GPU, B_GPU, C_GPU, ayDim, cyDim);

	err = cudaPeekAtLastError();
	if (err == cudaSuccess) { printf("\nThere was no error before synchronization!\n");}

	err = cudaDeviceSynchronize();
	if (err == cudaSuccess) { printf("\nNo problems synchronizing device!\n");}

	err = cudaPeekAtLastError();
	if (err == cudaSuccess) { printf("\nThere was no error after synchronization!\n");}

	err = cudaMemcpy(_C, C_GPU, axDim*byDim*sizeof(float), cudaMemcpyDeviceToHost);

	if (err == cudaSuccess) { printf("\nNo problems copying arrays from device!\n");}

	cudaFree(A_GPU);
	cudaFree(B_GPU);
	cudaFree(C_GPU);
	cudaDeviceReset();		
	
	printf("\nA\n");
	displayMatrix(A, ax, ay);

	printf("\nB\n");
	displayMatrix(B, bx, by);

	printf("\nC\n");
	displayMatrix(C, cx, cy);	

	printf("\n_A\n");
	displayMatrix(_A, axDim, ayDim);

	printf("\n_B\n");
	displayMatrix(_B, bxDim, byDim);

	printf("\n_C\n");
	displayMatrix(_C, cxDim, cyDim);

	//printf("\nWhere the fuck is the seg fault????");
	return 0;
}

float* fillMatrix(const int numRows, const int numCols, const float value)
{
	float* matrix = (float *)malloc(numRows*numCols*sizeof(float));	
	int i,j;

	for(i = 0; i < numRows; i++)
	{
		for(j = 0; j < numCols; j++)
		{
			matrix[(numCols*i)+j] = value;
		}
	}
	
	return matrix;
}

float* padMatrix(float* originalMatrix, const int numRows, const int numCols, const int newDim)
{
	float* matrix;
	int rows;
	int cols;

	if (numRows%newDim != 0) { rows = numRows - (numRows%newDim) + newDim; }
	if (numCols%newDim != 0) { cols = numCols - (numCols%newDim) + newDim; }

	matrix = (float *)malloc(rows*cols*sizeof(float));
	
	for (int i = 0; i < rows; i++)
	{
		for (int j = 0; j < cols; j++)
		{
			if (i < rows && j < cols)
			{
				matrix[cols*i+j] = originalMatrix[numCols*i+j];
			}
			else
			{
				matrix[cols*i+j] = 0.0f;
			}
		}	
	}

	return matrix;
}

void displayMatrix(const float* matrix, const int numRows, const int numCols)
{
	for (int i = 0; i < numRows; i++)
	{
		for ( int j = 0; j < numCols; j++)
		{
			printf(" %.2f ", matrix[numCols*i+j]);
		}
		printf("\n");
	}
}
